/*
*         (C) COPYRIGHT NBI Limited
*              ALL RIGHT RESERVED
*
* File        : fastflow/graph/schema_registry.h
* Authors     : dzhang
* Create Time : 2021/09/05 15:23:03
* Description :
*
*/

#ifndef __FASTFLOW_GRAPH_SCHEMA_REGISTRY_H___
#define __FASTFLOW_GRAPH_SCHEMA_REGISTRY_H___

#include "fastflow/graph/constants.h"
#include "fastflow/core/status.h"
#ifdef __GNUC__
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wignored-qualifiers"
#pragma GCC diagnostic ignored "-Wunused-parameter"
#endif
#include "onnx/defs/schema.h"
#ifdef __GNUC__
#pragma GCC diagnostic pop
#endif
#include <mutex>
#include <deque>
#include "sstream"

namespace fastflow {
using OpName_Domain_Version_Schema_Map = std::unordered_map<
    std::string,
    std::unordered_map<std::string, std::map<ONNX_NAMESPACE::OperatorSetVersion, ONNX_NAMESPACE::OpSchema>>>;

/** 
@struct SchemaRegistryVersion
fastflow schema registry is a supplement to the built-in ONNX schema.
Every schema registry represent a collection of schema deltas from baseline_opset_version to opset_version
*/
struct SchemaRegistryVersion {
  int baseline_opset_version;
  int opset_version;
};

using DomainToVersionMap = std::unordered_map<std::string, int>;
using DomainToVersionRangeMap = std::unordered_map<std::string, SchemaRegistryVersion>;

class IFastFlowOpSchemaCollection : public ONNX_NAMESPACE::ISchemaRegistry {
 public:
  virtual DomainToVersionMap GetLatestOpsetVersions(bool is_onnx_only) const = 0;

  using ISchemaRegistry::GetSchema;

  virtual const ONNX_NAMESPACE::OpSchema* GetSchema(
      const std::string& key,
      const int maxInclusiveVersion,
      const std::string& domain) const final {
    const ONNX_NAMESPACE::OpSchema* latest_schema = nullptr;
    int earliest_opset_where_unchanged = std::numeric_limits<int>::max();
    GetSchemaAndHistory(key, maxInclusiveVersion, domain, &latest_schema, &earliest_opset_where_unchanged);

    assert(latest_schema == nullptr || (latest_schema->SinceVersion() <= maxInclusiveVersion &&
                                        earliest_opset_where_unchanged == latest_schema->SinceVersion()));

    return latest_schema;
  }

  virtual void GetSchemaAndHistory(
      const std::string& key,
      int maxInclusiveVersion,
      const std::string& domain,
      const ONNX_NAMESPACE::OpSchema** latest_schema,
      int* earliest_opset_where_unchanged) const = 0;
};

/** 
@class FastFlowOpSchemaRegistry

FastFlowOpSchemaRegistry is used to provide supplement for built-in ONNX schemas.
Each FastFlowOpSchemaRegistry must register complete opsets delta from a baseline version to max opset version.
(Please notice that baseline opsets are not include in the delta)

For example, FASTFLOW is build with ONNX 1.2 which is at opset7, to use ONNX opset8 and opset9,
user could create a FastFlowOpSchemaRegistry and config it as {baseline_opset_version = 7, opset_version = 9}
it means this FastFlowOpSchemaRegistry contains the complete delta from opset7 to opset9.
*/
class FastFlowOpSchemaRegistry : public IFastFlowOpSchemaCollection {
 public:
  FastFlowOpSchemaRegistry() = default;

  Status SetBaselineAndOpsetVersionForDomain(
      const std::string& domain,
      int baseline_opset_version,
      int opset_version);

  DomainToVersionMap GetLatestOpsetVersions(bool is_onnx_only) const override;

  // FastFlowOpSchemaRegistry must register complete delta for a opset.
  Status RegisterOpSet(
      std::vector<ONNX_NAMESPACE::OpSchema>& schemas,
      const std::string& domain,
      int baseline_opset_version,
      int opset_version);

  using IFastFlowOpSchemaCollection::GetSchema;

  void GetSchemaAndHistory(
      const std::string& key,
      const int maxInclusiveVersion,
      const std::string& domain,
      const ONNX_NAMESPACE::OpSchema** latest_schema,
      int* earliest_opset_where_unchanged) const override;

  bool empty() const {
    return map_.empty();
  }

 private:
  Status RegisterOpSchema(ONNX_NAMESPACE::OpSchema&& op_schema);

  Status RegisterOpSchemaInternal(ONNX_NAMESPACE::OpSchema&& op_schema);

  std::mutex mutex_;

  OpName_Domain_Version_Schema_Map map_;
  DomainToVersionRangeMap domain_version_range_map_;
};

/**
@class SchemaRegistryManager 

SchemaRegistryManager provides a view based on built-in ONNX schema and a list of 
FastFlowOpSchemaRegistry as supplement.

The user needs to make sure the customized schema registry is valid, otherwise the behavior is undefined.

@todo We may add more consistency checks later.
*/
class SchemaRegistryManager : public fastflow::IFastFlowOpSchemaCollection {
 public:
  /**
  Register a new schema registry instance.
  @remarks The schema registry priority is the reverse of registration order. i.e. the last registry added will be
  searched first for a matching OpSchema.
  */
  void RegisterRegistry(std::shared_ptr<IFastFlowOpSchemaCollection> registry);

  /** Gets the latest opset versions.
  @param is_onnx_only If true, return the latest ONNX schemas. If false, return the latest schemas for all domains.
  */
  DomainToVersionMap GetLatestOpsetVersions(bool is_onnx_only) const override;

  /**
  Gets the OpSchema and its history.
  Searches custom schema registries starting with the last one added. \
  If the OpSchema is not found the default ONNX schema registry is searched.

  @param key Operator type.
  @param max_inclusive_version Maximum opset version allowed, inclusive.
  @param domain The domain of the operator.
  @param[out] latest_schema Returns the latest OpSchema if found. nullptr otherwise.
  @param[out] earliest_opset_where_unchanged The earliest opset version preceding max_inclusive_version where the
  operator is known to be unchanged.
  */
  void GetSchemaAndHistory(
      const std::string& key,
      const int max_inclusive_version,
      const std::string& domain,
      const ONNX_NAMESPACE::OpSchema** latest_schema,
      int* earliest_opset_where_unchanged) const override;

 private:
  std::deque<std::shared_ptr<IFastFlowOpSchemaCollection>> registries;
};

}  // namespace fastflow

#endif // __FASTFLOW_GRAPH_SCHEMA_REGISTRY_H___
